options(java.parameters = "-Xmx50000m")
library(rnn)
library(grf)
library(bartMachine)
library(doParallel)
library(foreach)
library(tictoc)
library(glmnet)
source("generate_testing_data.R")

PAPE <- function (T, That, Y) {
  n=length(Y)
  n1=sum(T)
  n0=n-n1
  n1h=sum(That)
  n0h=n-n1h
  probs=sum(That)/n
  SAPE=n/(n-1)*(1/n1*sum(T*That*Y)+1/n0*sum(Y*(1-T)*(1-That))-n1h/(n1*n)*sum(Y*T)-n0h/(n0*n)*sum(Y*(1-T)))
  Sf1=var(((That-probs)*Y)[T==1])
  Sf0=var(((That-probs)*Y)[T==0])
  SATE=1/n1*sum(T*Y)-1/n0*(sum((1-T)*Y))
  covarterm=1/n^2*(SAPE^2+2*(n-1)*SAPE*SATE*(2*probs-1)-(1-probs)*probs*n*SATE^2)
  varexp=n^2/(n-1)^2*(Sf1/n1+Sf0/n0+covarterm)
  return(list(pape=SAPE,sd=sqrt(varexp)))
}

PAPEp <- function (T, That, Y, plim) {
  n=length(Y)
  n1=sum(T)
  n0=n-n1
  n1h=sum(That)
  n0h=n-n1h
  SAPEfp=1/n1*sum(T*That*Y)+1/n0*sum(Y*(1-T)*(1-That))-plim/n1*sum(Y*T)-(1-plim)/n0*sum(Y*(1-T))
  Sfp1=var(((That-plim)*Y)[T==1])
  Sfp0=var(((That-plim)*Y)[T==0])
  kf1=mean(Y[T==1 & That==1])-mean(Y[T==0 & That==1])
  kf0=mean(Y[T==1 & That==0])-mean(Y[T==0 & That==0])
  varfp=Sfp1/n1+Sfp0/n0+floor(n*plim)*(n-floor(n*plim))/(n^2*(n-1))*((2*plim-1)*kf1^2-2*plim*kf1*kf0)
  return(list(papep=SAPEfp,sd=sqrt(varfp)))
}

PAPD <- function (T, Thatfp, Thatgp, Y, plim) {
  n=length(Y)
  n1=sum(T)
  n0=n-n1
  SAPEfp=1/n1*sum(T*Thatfp*Y)+1/n0*sum(Y*(1-T)*(1-Thatfp))-plim/n1*sum(Y*T)-(1-plim)/n0*sum(Y*(1-T))
  SAPEgp=1/n1*sum(T*Thatgp*Y)+1/n0*sum(Y*(1-T)*(1-Thatgp))-plim/n1*sum(Y*T)-(1-plim)/n0*sum(Y*(1-T))
  Sfp1=var(((Thatfp-plim)*Y)[T==1])
  Sfp0=var(((Thatfp-plim)*Y)[T==0])
  kf1=mean(Y[T==1 & Thatfp==1])-mean(Y[T==0 & Thatfp==1])
  kf0=mean(Y[T==1 & Thatfp==0])-mean(Y[T==0 & Thatfp==0])
  PAPD=SAPEfp-SAPEgp
  Sfgp1=var(((Thatfp-Thatgp)*Y)[T==1])
  Sfgp0=var(((Thatfp-Thatgp)*Y)[T==0])
  kg1=mean(Y[T==1 & Thatgp==1])-mean(Y[T==0 & Thatgp==1])
  kg0=mean(Y[T==1 & Thatgp==0])-mean(Y[T==0 & Thatgp==0])
  varfgp=Sfgp1/n1+Sfgp0/n0-floor(n*plim)*(n-floor(n*plim))/(n^2*(n-1))*(kf1^2+kg1^2)+
    2*floor(n*plim)*max(floor(n*plim),n-floor(n*plim))/(n^2*(n-1))*abs(kf1*kg1)
  if (varfgp>0) {
    return(list(papd=PAPD,sd=sqrt(varfgp)))
  } else {
    return(list(papd=PAPD,sd=0))
  }
}


AUPEC <- function (T, tau, Y) {
  n=length(Y)
  n1=sum(T)
  n0=n-n1
  pf=sum(as.numeric(tau>0))/n
  Z=rbinom(1e4,n,pf)
  
  ThatfA=numeric(n)
  kAf1=numeric(n)
  kAf0=numeric(n)
  kAf1A=numeric(n)
  kAf1B=numeric(n)
  kAf0A=numeric(n)
  kAf0B=numeric(n)
  covarsum=0
  for (i in 1:n) {
    cutofftemp=quantile(tau,1-i/n)
    Thatftemp=as.numeric(tau>cutofftemp)
    ThatftempA=as.numeric(tau>max(cutofftemp,0))
    ThatfA=ThatfA+1/n*ThatftempA
    cutofftemp2=quantile(tau,(i-1)/n)
    Thatftemp2=as.numeric(tau>cutofftemp2)
    tempkAf1=mean(Y[T==1 & Thatftemp2==1])-mean(Y[T==0 & Thatftemp2==1])
    tempkAf0=mean(Y[T==1 & Thatftemp==0])-mean(Y[T==0 & Thatftemp==0])
    if (is.nan(tempkAf1)) {
      kAf1[n-i+1]=kAf1[n-i+2]
    } else {
      kAf1[n-i+1]=tempkAf1
    }
    if (is.nan(tempkAf0)) {
      kAf0[i]=kAf0[i-1]
    } else {
      kAf0[i]=tempkAf0
    }
    kAf1A[n-i+1]=(n-i+1)*kAf1[n-i+1]
    kAf1B[n-i+1]=(i-1)*kAf1[n-i+1]
    kAf0A[i]=i*kAf0[i]
    kAf0B[i]=(n-i)*kAf0[i]
  }
  sumtemp1=cumsum(kAf1A*kAf0B)
  sumtemp2=cumsum(kAf1A)
  sumtemp3=cumsum(kAf1A*kAf1B)
  tempM=outer(kAf1A,kAf1B)
  tempM[lower.tri(tempM,diag=TRUE)] <- 0
  tempMsum=cumsum(colSums(tempM))
  covarsum1=mean(-1/(n^3*(n-1))*sumtemp1[Z]-Z*(n-Z)^2/(n^3*(n-1))*kAf1[Z]*kAf0[Z]-2/(n^4*(n-1))*tempMsum[Z]-Z^2*(n-Z)^2/(n^4*(n-1))*kAf1[Z]^2
                 -2*(n-Z)^2/(n^4*(n-1))*kAf1[Z]*sumtemp2[Z]+1/n^4*sumtemp3[Z])
  covarsum2=var(1/n*(sumtemp2[Z]/n+(n-Z)*Z/n*kAf1[Z]))
  ThatfA2=ThatfA-1/2
  SfA1=var((ThatfA2*Y)[T==1])
  SfA0=var((ThatfA2*Y)[T==0])
  varfA=SfA1/n1+SfA0/n0+covarsum1+covarsum2
  AUPEC=1/n1*sum(T*ThatfA*Y)+1/n0*sum(Y*(1-T)*(1-ThatfA))-0.5/n1*sum(T*Y)-0.5/n0*sum((1-T)*Y)
  return(list(aupec=AUPEC,sd=sqrt(varfA)))
}




ntrials=1000



### PAPE Experiments
no_cores <- detectCores() -1
cl <- makeCluster(no_cores)
registerDoParallel(cl)
final_output=matrix(0, 6, 10)
m=1
for (i in 0:1) {
  magnitude = i
  foldername = paste(magnitude,sep='')
  train = read.csv(paste('iid/',foldername,'.csv',sep=''))
  y_train = train$y
  xz_train = train[, !names(train) %in% c("y")]
  x_train = train[, !names(train) %in% c("y","z")]
  x_train_expand = model.matrix(~.-1,data=x_train)
  z_train = train$z
  cf1<-causal_forest(x_train_expand,y_train,z_train,tune.parameters = TRUE,num.trees = 4000)
  nseq=c(100,500,2000)
  x = read.csv('Xt.csv',header=FALSE)
  x_full = x[,c(1,3,10,14,15,21,24,43)]
  x_full_expand = model.matrix(~.-1,data=x_full)
  tau_full1<-predict(cf1,x_full_expand)
  That = as.numeric(tau_full1$predictions>0) 
  PAPEavg<- truePAPE(That,magnitude)
  for (k in 1:length(nseq)) {
    PAPEu = numeric(ntrials)
    PAPEl = numeric(ntrials)
    PAPEm = numeric(ntrials)
    sdm = numeric(ntrials)
    tic()
    expresult=foreach (l=1:ntrials) %dopar% {
      library(grf)
      test = testData(nseq[k],magnitude)
      y_test = test$y
      xz_test = test[, !names(test) %in% c("y")]
      x_test = test[, !names(test) %in% c("y","z")]
      z_test = test$z
      x_test_expand = model.matrix(~.-1,data=x_test)
      tau_test1<-predict(cf1,x_test_expand)
      That = as.numeric(tau_test1$predictions>0) 
      output = PAPE(z_test,That,y_test)
      pape = output$pape
      sdt = output$sd
      return(list(pape+1.96*sdt, pape, pape-1.96*sdt, sdt))
    }
    toc()
    PAPEu=unlist(expresult)[c(TRUE,FALSE,FALSE,FALSE)]
    PAPEm=unlist(expresult)[c(FALSE,TRUE,FALSE,FALSE)]
    PAPEl=unlist(expresult)[c(FALSE,FALSE,TRUE,FALSE)]
    sdm=unlist(expresult)[c(FALSE,FALSE,FALSE,TRUE)]
    coverage=sum(as.numeric(PAPEavg<PAPEu & PAPEavg>PAPEl))/ntrials
    rmse = sqrt(mean((PAPEm-PAPEavg)*(PAPEm-PAPEavg)))
    mad = mean(abs(PAPEm-PAPEavg))
    asd=sd(PAPEm)
    rmsesd= sqrt(mean((sdm-asd)*(sdm-asd)))
    madsd = mean(abs(sdm-asd))
    final_output[m,1]=magnitude*10
    final_output[m,2]=nseq[k]
    final_output[m,3]=coverage
    final_output[m,4]=PAPEavg
    final_output[m,5]=mean(PAPEm)-PAPEavg
    final_output[m,6]=asd
    final_output[m,7]=rmse
    final_output[m,8]=mad
    final_output[m,9]=rmsesd
    final_output[m,10]=madsd
    m = m+1
  }
}
stopCluster(cl)


plim = 0.2
## PAPEp Experiments
no_cores <- detectCores() -1
cl <- makeCluster(no_cores)
registerDoParallel(cl)
final_output2=matrix(0, 6, 10)
m=1
for (i in 0:1) {
  magnitude = i
  foldername = paste(magnitude,sep='')
  train = read.csv(paste('iid/',foldername,'.csv',sep=''))
  y_train = train$y
  xz_train = train[, !names(train) %in% c("y")]
  x_train = train[, !names(train) %in% c("y","z")]
  x_train_expand = model.matrix(~.-1,data=x_train)
  z_train = train$z
  cf1<-causal_forest(x_train_expand,y_train,z_train,tune.parameters = TRUE,num.trees = 4000)
  nseq=c(100,500,2000)
  x = read.csv('Xt.csv',header=FALSE)
  x_full = x[,c(1,3,10,14,15,21,24,43)]
  x_full_expand = model.matrix(~.-1,data=x_full)
  tau_full1<-predict(cf1,x_full_expand)
  That = as.numeric(tau_full1$predictions>quantile(tau_full1$predictions,1-plim))
  PAPEavg<- truePAPEp(That,magnitude,plim)
  for (k in 1:length(nseq)) {
    PAPEu = numeric(ntrials)
    PAPEl = numeric(ntrials)
    PAPEm = numeric(ntrials)
    sdm = numeric(ntrials)
    tic()
    expresult=foreach (l=1:ntrials) %dopar% {
      library(grf)
      test = testData(nseq[k],magnitude)
      y_test = test$y
      xz_test = test[, !names(test) %in% c("y")]
      x_test = test[, !names(test) %in% c("y","z")]
      z_test = test$z
      x_test_expand = model.matrix(~.-1,data=x_test)
      tau_test1<-predict(cf1,x_test_expand)
      That = as.numeric(tau_test1$predictions>quantile(tau_test1$predictions,1-plim))
      output = PAPEp(z_test,That,y_test,plim)
      pape = output$pape
      sdt = output$sd
      return(list(pape+1.96*sdt, pape, pape-1.96*sdt, sdt))
    }
    toc()
    PAPEu=unlist(expresult)[c(TRUE,FALSE,FALSE,FALSE)]
    PAPEm=unlist(expresult)[c(FALSE,TRUE,FALSE,FALSE)]
    PAPEl=unlist(expresult)[c(FALSE,FALSE,TRUE,FALSE)]
    sdm=unlist(expresult)[c(FALSE,FALSE,FALSE,TRUE)]
    coverage=sum(as.numeric(PAPEavg<PAPEu & PAPEavg>PAPEl))/ntrials
    rmse = sqrt(mean((PAPEm-PAPEavg)*(PAPEm-PAPEavg)))
    mad = mean(abs(PAPEm-PAPEavg))
    asd=sd(PAPEm)
    rmsesd= sqrt(mean((sdm-asd)*(sdm-asd)))
    madsd = mean(abs(sdm-asd))
    final_output2[m,1]=magnitude*10
    final_output2[m,2]=nseq[k]
    final_output2[m,3]=coverage
    final_output2[m,4]=PAPEavg
    final_output2[m,5]=mean(PAPEm)-PAPEavg
    final_output2[m,6]=asd
    final_output2[m,7]=rmse
    final_output2[m,8]=mad
    final_output2[m,9]=rmsesd
    final_output2[m,10]=madsd
    m = m+1
  }
}
stopCluster(cl)


# PAPD experiments (LASSO)
no_cores <- detectCores() -1
cl <- makeCluster(no_cores)
registerDoParallel(cl)

final_output3b=matrix(0, 6, 10)
m=1
plim=0.2
for (i in 0:1) {
  magnitude = i
  foldername = paste(magnitude,sep='')
  train = read.csv(paste('iid/',foldername,'.csv',sep=''))
  y_train = train$y
  xz_train = train[, !names(train) %in% c("y")]
  xz_train_expand<-model.matrix(~.*z,data=xz_train)
  x_train_expand = model.matrix(~.-1,data=x_train)
  z_train = train$z
  cf1<-causal_forest(x_train_expand,y_train,z_train,tune.parameters = TRUE,num.trees = 4000)
  ls1<-glmnet(xz_train_expand,y_train,alpha=1,lambda=0.05)
  x = read.csv('Xt.csv',header=FALSE)
  x_full = x[,c(1,3,10,14,15,21,24,43)]
  x_full0 = cbind(x_full, z=0)
  x_full1 = cbind(x_full, z=1)
  x_full0_expand=model.matrix(~.*z,data=x_full0)
  x_full1_expand=model.matrix(~.*z,data=x_full1)
  x_full_expand = model.matrix(~.-1,data=x_full)
  yl0 = predict(ls1, x_full0_expand)
  yl1 = predict(ls1, x_full1_expand)  
  tau_full0<-predict(cf1,x_full_expand)$predictions
  tau_full1<-yl1-yl0
  Thatfp = numeric(length(tau_full0))
  Thatfp[sort(tau_full0,decreasing = TRUE,index.return=TRUE)$ix[1:(round(plim*length(Thatfp))+1)]] = 1
  Thatgp = numeric(length(Thatfp))
  Thatgp[sort(tau_full1,decreasing = TRUE,index.return=TRUE)$ix[1:(round(plim*length(Thatfp))+1)]] = 1
  PAPDavg<- truePAPD(Thatfp,Thatgp, magnitude,plim)
  nseq=c(100,500,2000)
  for (k in 1:length(nseq)) {
    tic()
    expresult=foreach (l=1:ntrials) %dopar% {
      library(grf)
      library(glmnet)
      test = testData(nseq[k],magnitude)
      y_test = test$y
      x_test = test[, !names(test) %in% c("y","z")]
      x_test0 = cbind(x_test, z=0)
      x_test1 = cbind(x_test, z=1)
      x_test0_expand=model.matrix(~.*z,data=x_test0)
      x_test1_expand=model.matrix(~.*z,data=x_test1)
      x_test_expand = model.matrix(~.-1,data=x_test)
      z_test = test$z
      tau_test0<-predict(cf1,x_test_expand)$predictions
      Thatfp = numeric(length(tau_test0))
      Thatfp[sort(tau_test0,decreasing = TRUE,index.return=TRUE)$ix[1:(round(plim*length(Thatfp))+1)]] = 1
      yl0 <-predict(ls1,x_test0_expand)
      yl1 <-predict(ls1,x_test1_expand)
      tau_test1<-yl1-yl0
      Thatgp = numeric(length(Thatfp))
      Thatgp[sort(tau_test1,decreasing = TRUE,index.return=TRUE)$ix[1:(round(plim*length(Thatfp))+1)]] = 1
      output = PAPD(z_test,Thatfp,Thatgp,y_test,plim)
      papd = output$papd
      sdt = output$sd
      return(list(papd+1.96*sdt, papd, papd-1.96*sdt, sdt))
    }
    toc()
    PAPDu=unlist(expresult)[c(TRUE,FALSE,FALSE,FALSE)]
    PAPDm=unlist(expresult)[c(FALSE,TRUE,FALSE,FALSE)]
    PAPDl=unlist(expresult)[c(FALSE,FALSE,TRUE,FALSE)]
    sdm=unlist(expresult)[c(FALSE,FALSE,FALSE,TRUE)]
    coverage=sum(as.numeric(PAPDavg<PAPDu & PAPDavg>PAPDl))/ntrials
    rmse = sqrt(mean((PAPDm-PAPDavg)*(PAPDm-PAPDavg)))
    mad = mean(abs(PAPDm-PAPDavg))
    asd=sd(PAPDm)
    rmsesd= sqrt(mean((sdm-asd)*(sdm-asd)))
    madsd = mean(abs(sdm-asd))
    final_output3b[m,1]=magnitude*10
    final_output3b[m,2]=nseq[k]
    final_output3b[m,3]=coverage
    final_output3b[m,4]=PAPDavg
    final_output3b[m,5]=mean(PAPDm)-PAPDavg
    final_output3b[m,6]=asd
    final_output3b[m,7]=rmse
    final_output3b[m,8]=mad
    final_output3b[m,9]=rmsesd
    final_output3b[m,10]=madsd
    m = m+1
  }
}
stopCluster(cl)

no_cores <- detectCores() -1
cl <- makeCluster(no_cores)
registerDoParallel(cl)

### AUPEC Experiments
final_output4=matrix(0, 6, 10)
m=1
for (i in 0:1) {
  magnitude = i
  foldername = paste(magnitude,sep='')
  train = read.csv(paste('iid/',foldername,'.csv',sep=''))
  y_train = train$y
  xz_train = train[, !names(train) %in% c("y")]
  x_train = train[, !names(train) %in% c("y","z")]
  x_train_expand = model.matrix(~.-1,data=x_train)
  z_train = train$z
  cf1<-causal_forest(x_train_expand,y_train,z_train,tune.parameters = TRUE,num.trees = 4000)
  nseq=c(100,500,2000)
  x = read.csv('Xt.csv',header=FALSE)
  x_full = x[,c(1,3,10,14,15,21,24,43)]
  x_full_expand = model.matrix(~.-1,data=x_full)
  tau_full0<-predict(cf1,x_full_expand)
  AUPECavg<-trueAUPEC(tau_full0$predictions,magnitude)
  print(AUPECavg)
  nseq=c(100,500,2000)
  for (k in 1:length(nseq)) {
    tic()
    expresult=foreach (l=1:ntrials) %dopar% {
      library(grf)
      test = testData(nseq[k],magnitude)
      y_test = test$y
      xz_test = test[, !names(test) %in% c("y")]
      x_test = test[, !names(test) %in% c("y","z")]
      z_test = test$z
      x_test_expand = model.matrix(~.-1,data=x_test)
      tau_test1<-predict(cf1,x_test_expand)
      output = AUPEC(z_test,tau_test1$predictions,y_test)
      aupec = output$aupec
      sdt = output$sd
      return(list(aupec+1.96*sdt, aupec, aupec-1.96*sdt, sdt))
    }
    toc()
    AUPECu=unlist(expresult)[c(TRUE,FALSE,FALSE,FALSE)]
    AUPECm=unlist(expresult)[c(FALSE,TRUE,FALSE,FALSE)]
    AUPECl=unlist(expresult)[c(FALSE,FALSE,TRUE,FALSE)]
    sdm=unlist(expresult)[c(FALSE,FALSE,FALSE,TRUE)]
    coverage=sum(as.numeric(AUPECavg<AUPECu & AUPECavg>AUPECl))/ntrials
    rmse = sqrt(mean((AUPECm-AUPECavg)*(AUPECm-AUPECavg)))
    mad = mean(abs(AUPECm-AUPECavg))
    asd=sd(AUPECm)
    rmsesd= sqrt(mean((sdm-asd)*(sdm-asd)))
    madsd = mean(abs(sdm-asd))
    final_output4[m,1]=magnitude*10
    final_output4[m,2]=nseq[k]
    final_output4[m,3]=coverage
    final_output4[m,4]=mean(AUPECm)-AUPECavg
    final_output4[m,5]=asd
    final_output4[m,6]=rmse
    final_output4[m,7]=mad
    final_output4[m,8]=rmsesd
    final_output4[m,9]=madsd
    m = m+1
  }
}
stopCluster(cl)